import sys
import shutil
import subprocess
import tempfile
import ctypes
import time
import hashlib
import zipfile
import os
import signal
from pkg_resources import parse_version
from pathlib import Path
from typing import Tuple, Optional, List
import threading
import requests
import customtkinter as ctk
from queue import Queue, Empty
from tkinter import messagebox
import pefile  # 需安装：pip install pefile

# ================== 配置区 ==================
UPDATE_CONFIG = {
    "version_url": "http://115.120.195.74/latest.json",  # 远程版本配置地址
    "timeout": 8,
}

# 路径初始化
APP_ROOT = Path(__file__).parent if not getattr(sys, 'frozen', False) else Path(sys.executable).parent
CURRENT_EXE = Path(sys.executable).resolve()

UPDATE_TEMP_DIR = APP_ROOT / "update_temp"
UPDATE_TEMP_DIR.mkdir(exist_ok=True)
TEMP_ZIP_FILE = UPDATE_TEMP_DIR / f"update_{CURRENT_EXE.stem}.zip"
TEMP_UNZIP_DIR = UPDATE_TEMP_DIR / f"update_{CURRENT_EXE.stem}_temp"

# ================== 日志初始化 ==================
try:
    from logger import app_logger as logger
except ImportError:
    import logging
    log_dir = APP_ROOT / "logs"
    log_dir.mkdir(exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_dir / "app.log", encoding='utf-8'),
            logging.StreamHandler(sys.stdout)
        ]
    )
    logger = logging.getLogger("App")

# ================== 共享工具函数 ==================
def verify_sha256(file_path: Path, expected_hash: str) -> bool:
    if not file_path.exists():
        logger.error(f"[校验] 文件不存在：{file_path}")
        return False
    sha256 = hashlib.sha256()
    try:
        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha256.update(chunk)
        actual = sha256.hexdigest().lower()
        expected = expected_hash.lower()
        logger.info(f"[校验] 实际 SHA256: {actual}")
        logger.info(f"[校验] 期望 SHA256: {expected}")
        return actual == expected
    except Exception as e:
        logger.error(f"[校验] 校验失败: {e}", exc_info=True)
        return False

def is_64bit_file(file_path: Path) -> Optional[bool]:
    logger.info(f"[架构识别] 检测文件架构：{file_path.name}（路径：{file_path}）")
    if not file_path.exists() or not file_path.is_file():
        logger.error(f"[架构识别] 文件不存在或不是文件：{file_path}")
        return None
    try:
        pe = pefile.PE(str(file_path))
        machine_type = pe.FILE_HEADER.Machine
        if machine_type == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]:
            logger.info(f"[架构识别] 识别为64位文件")
            return True
        elif machine_type == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]:
            logger.info(f"[架构识别] 识别为32位文件")
            return False
        else:
            logger.warning(f"[架构识别] 未知架构（Machine类型：{machine_type}）")
            return None
    except pefile.PEFormatError as e:
        logger.error(f"[架构识别] 非PE格式文件：{e}")
        return None
    except Exception as e:
        logger.error(f"[架构识别] 解析失败：{e}", exc_info=True)
        return None

def is_64bit_os() -> bool:
    result = sys.maxsize > 2 ** 32
    logger.info(f"[系统信息] 检测到{'64位' if result else '32位'}操作系统")
    return result

def is_admin() -> bool:
    try:
        result = ctypes.windll.shell32.IsUserAnAdmin()
        logger.info(f"[权限检查] 管理员权限：{'是' if result else '否'}")
        return result
    except:
        logger.warning("[权限检查] 权限检查失败，默认返回非管理员")
        return False

def is_process_running(exe_name: str) -> bool:
    try:
        import psutil
        for proc in psutil.process_iter(['name']):
            if proc.info['name'] == exe_name:
                return True
        return False
    except ImportError:
        output = subprocess.check_output(f'tasklist /FI "IMAGENAME eq {exe_name}"', shell=True, text=True)
        return exe_name in output
    except Exception as e:
        logger.warning(f"检查进程失败: {e}")
        return False

def kill_process_by_name(process_name: str) -> bool:
    try:
        import psutil
        current_pid = os.getpid()
        killed = False
        for proc in psutil.process_iter(['pid', 'name']):
            if proc.info['name'] == process_name and proc.info['pid'] != current_pid:
                proc.kill()
                proc.wait(timeout=3)
                logger.info(f"[退出流程] 已强制终止同名进程 PID {proc.info['pid']}")
                killed = True
        return killed
    except ImportError:
        logger.warning("[退出流程] psutil 未安装，跳过同名进程清理")
        return False
    except Exception as e:
        logger.error(f"[退出流程] kill_process_by_name 异常: {e}")
        return False

def clean_old_mei_tempdirs():
    try:
        temp_dir = Path(tempfile.gettempdir())
        count = 0
        for item in temp_dir.iterdir():
            if item.is_dir() and item.name.startswith("_MEI"):
                if hasattr(sys, '_MEIPASS') and str(item) == sys._MEIPASS:
                    continue
                try:
                    shutil.rmtree(item)
                    count += 1
                except PermissionError:
                    pass
        logger.info(f"[清理] 临时目录清理完成，共删除 {count} 个目录")
    except Exception as e:
        logger.warning(f"[清理] 清理临时目录异常: {e}")

def get_current_version() -> str:
    ver_file = APP_ROOT / "version.txt"
    if ver_file.exists():
        try:
            with open(ver_file, 'r', encoding='utf-8') as f:
                return f.read().strip()
        except Exception as e:
            logger.error(f"[版本获取] 读取 {ver_file} 失败: {e}")
    return "1.0.0"

def get_remote_version() -> Tuple[Optional[str], Optional[str], str, Optional[str], Optional[str]]:
    try:
        resp = requests.get(UPDATE_CONFIG["version_url"], timeout=UPDATE_CONFIG["timeout"])
        resp.raise_for_status()
        data = resp.json()
        return (
            data.get("version"),
            data.get("url"),
            data.get("changelog", ""),
            data.get("exe_sha256"),
            data.get("arch")
        )
    except Exception as e:
        logger.error(f"[远程版本] 获取失败: {e}", exc_info=True)
        return None, None, "", None, None

def show_update_dialog(current_ver: str, remote_ver: str, changelog: str) -> bool:
    root = ctk.CTk()
    root.title("发现新版本")
    root.geometry("520x460")
    root.resizable(False, False)
    root.attributes("-topmost", True)
    root.protocol("WM_DELETE_WINDOW", lambda: None)

    main_frame = ctk.CTkFrame(root)
    main_frame.pack(fill="both", expand=True, padx=15, pady=15)

    ctk.CTkLabel(main_frame, text="🎉 发现新版本！", font=("Microsoft YaHei", 18, "bold")).pack(pady=(10, 5))
    ctk.CTkLabel(main_frame, text=f"当前版本：{current_ver} → 新版本：{remote_ver}").pack()

    ctk.CTkLabel(main_frame, text="更新内容：", font=("Microsoft YaHei", 12, "underline"), anchor="w").pack(anchor="w", padx=10, pady=(15, 5))
    text_frame = ctk.CTkFrame(main_frame, fg_color="transparent")
    text_frame.pack(fill="both", expand=True, padx=10, pady=(0, 15))
    textbox = ctk.CTkTextbox(text_frame, wrap="word", font=("Microsoft YaHei", 11))
    textbox.pack(side="left", fill="both", expand=True)
    textbox.insert("0.0", changelog if changelog else "暂无更新说明。")
    textbox.configure(state="disabled")
    scrollbar = ctk.CTkScrollbar(text_frame, command=textbox.yview)
    scrollbar.pack(side="right", fill="y")
    textbox.configure(yscrollcommand=scrollbar.set)

    result = {"update_now": False}
    def on_update():
        result["update_now"] = True
        root.destroy()
    def on_later():
        result["update_now"] = False
        root.destroy()

    button_frame = ctk.CTkFrame(main_frame, fg_color="transparent")
    button_frame.pack(fill="x", pady=(0, 10))
    ctk.CTkButton(button_frame, text="立即更新", command=on_update, fg_color="green", width=120).pack(side="left", padx=(0, 15))
    ctk.CTkButton(button_frame, text="稍后再说", command=on_later, fg_color="gray", width=120).pack(side="right")

    root.grab_set()
    root.mainloop()
    return result["update_now"]

def extract_zip(zip_path: Path, extract_dir: Path) -> List[Path]:
    if extract_dir.exists():
        shutil.rmtree(extract_dir)
    extract_dir.mkdir(exist_ok=True)
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_dir)
        extracted_files = [extract_dir / f for f in zip_ref.namelist()]
    logger.info(f"[解压] 成功解压 {len(extracted_files)} 个文件到 {extract_dir}")
    return extracted_files

def find_required_files(extract_dir: Path):
    """只查找三个必需文件"""
    main_exe = extract_dir / "DataCollector.exe"
    updater_exe = extract_dir / "updater.exe"
    version_txt = extract_dir / "version.txt"
    return main_exe, updater_exe, version_txt

def perform_update_flow(current_ver: str, remote_ver: str, download_url: str, changelog: str, expected_sha256: str, remote_arch: str) -> bool:
    # 架构兼容性校验
    is_64 = is_64bit_os()
    if is_64 and remote_arch not in ("x64", "64"):
        messagebox.showerror("架构不兼容", f"当前系统为64位，但更新包为{remote_arch}位，无法更新！")
        return False
    if not is_64 and remote_arch not in ("x86", "32"):
        messagebox.showerror("架构不兼容", f"当前系统为32位，但更新包为{remote_arch}位，无法更新！")
        return False

    if not is_admin():
        messagebox.showerror("权限不足", "更新需要管理员权限。请右键程序，选择“以管理员身份运行”后再试。")
        return False

    updater_win = ctk.CTk()
    updater_win.title("正在更新...")
    updater_win.geometry("500x200")
    updater_win.resizable(False, False)
    updater_win.attributes("-topmost", True)
    status_label = ctk.CTkLabel(updater_win, text="准备中...", font=("Microsoft YaHei", 12))
    status_label.pack(pady=5)
    progress_bar = ctk.CTkProgressBar(updater_win, mode="determinate")
    progress_bar.set(0.0)
    progress_bar.pack(pady=5, padx=20, fill="x")
    speed_label = ctk.CTkLabel(updater_win, text="速度: -- KB/s", font=("Microsoft YaHei", 10), text_color="gray")
    speed_label.pack()
    cancel_event = threading.Event()
    result_queue = Queue()
    update_queue = Queue()

    def background_update_flow():
        try:
            if TEMP_ZIP_FILE.exists():
                TEMP_ZIP_FILE.unlink()
            if TEMP_UNZIP_DIR.exists():
                shutil.rmtree(TEMP_UNZIP_DIR)

            # 下载
            update_queue.put(("status", "正在连接服务器..."))
            resp = requests.get(download_url, stream=True, timeout=30)
            resp.raise_for_status()
            total = int(resp.headers.get('content-length', 0))
            downloaded = 0
            last_time = time.time()
            last_downloaded = 0
            with open(TEMP_ZIP_FILE, 'wb') as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    if cancel_event.is_set():
                        if TEMP_ZIP_FILE.exists():
                            TEMP_ZIP_FILE.unlink()
                        raise Exception("用户取消更新")
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        now = time.time()
                        elapsed = now - last_time
                        if elapsed > 0.5:
                            speed_kb = (downloaded - last_downloaded) / elapsed / 1024
                            speed_text = f"速度: {speed_kb:.1f} KB/s | 已下载: {downloaded // 1024} KB"
                            update_queue.put(("speed", speed_text))
                            last_time = now
                            last_downloaded = downloaded
                        if total > 0:
                            progress = min(0.7, downloaded / total * 0.7)
                            update_queue.put(("progress", progress))
            update_queue.put(("status", "正在解压更新包..."))
            extract_zip(TEMP_ZIP_FILE, TEMP_UNZIP_DIR)

            # 查找必需文件
            main_exe, updater_exe, version_txt = find_required_files(TEMP_UNZIP_DIR)
            if not all([main_exe.exists(), updater_exe.exists(), version_txt.exists()]):
                raise Exception("更新包缺少必需文件：DataCollector.exe、updater.exe 或 version.txt")

            # 校验主程序
            update_queue.put(("status", "正在校验主程序完整性..."))
            if not verify_sha256(main_exe, expected_sha256):
                raise Exception("主程序SHA256校验失败")

            # 校验架构
            update_queue.put(("status", "正在校验文件架构..."))
            is_64bit = is_64bit_file(main_exe)
            if is_64bit is None:
                raise Exception("无法识别主程序架构")
            if is_64bit_os() and not is_64bit:
                raise Exception("主程序为32位，与64位系统不兼容")
            if not is_64bit_os() and is_64bit:
                raise Exception("主程序为64位，与32位系统不兼容")

            update_queue.put(("status", "校验通过，准备更新..."))
            result_queue.put(("ready_to_update", main_exe, updater_exe, version_txt, expected_sha256))

        except Exception as e:
            error_msg = str(e)
            logger.error(f"[更新] 流程失败: {error_msg}", exc_info=True)
            result_queue.put(("error", error_msg))
        finally:
            if TEMP_ZIP_FILE.exists():
                try:
                    TEMP_ZIP_FILE.unlink()
                except:
                    pass

    thread = threading.Thread(target=background_update_flow, daemon=True)
    thread.start()

    def on_cancel():
        cancel_event.set()
        cancel_button.configure(state="disabled")
        status_label.configure(text="正在取消...")

    cancel_button = ctk.CTkButton(updater_win, text="取消更新", command=on_cancel, fg_color="red", width=100)
    cancel_button.pack(pady=10)

    def process_ui_updates():
        while True:
            try:
                msg_type, payload = update_queue.get_nowait()
                if msg_type == "status":
                    status_label.configure(text=payload)
                elif msg_type == "progress":
                    progress_bar.set(payload)
                elif msg_type == "speed":
                    speed_label.configure(text=payload)
            except Empty:
                break

        if not result_queue.empty():
            msg_type, *payload = result_queue.get()
            if msg_type == "ready_to_update":
                main_exe, updater_exe, version_txt, sha256 = payload
                updater_win.destroy()
                exit_and_launch_updater(main_exe, updater_exe, version_txt, sha256)
                return
            else:
                messagebox.showerror("更新失败", payload[0])
                updater_win.destroy()
                if TEMP_UNZIP_DIR.exists():
                    shutil.rmtree(TEMP_UNZIP_DIR, ignore_errors=True)
                return

        if thread.is_alive():
            updater_win.after(100, process_ui_updates)
        else:
            updater_win.destroy()

    updater_win.after(100, process_ui_updates)
    updater_win.mainloop()

def exit_and_launch_updater(main_exe: Path, updater_exe: Path, version_txt: Path, expected_sha256: str):
    logger.info("[退出流程] 主程序准备主动退出...")
    kill_process_by_name(CURRENT_EXE.name)
    for thread in threading.enumerate():
        if thread != threading.main_thread():
            try:
                thread.join(timeout=2.0)
            except:
                pass
    clean_old_mei_tempdirs()
    time.sleep(2)
    kill_process_by_name(CURRENT_EXE.name)
    time.sleep(1)

    try:
        if getattr(sys, 'frozen', False):
            updater_executable = APP_ROOT / "updater.exe"
            cmd = [
                f'"{updater_executable}"',
                f'"{main_exe}"',
                f'"{updater_exe}"',
                f'"{version_txt}"',
                expected_sha256
            ]
        else:
            updater_script = APP_ROOT / "updater.py"
            cmd = [
                f'"{sys.executable}"',
                f'"{updater_script}"',
                f'"{main_exe}"',
                f'"{updater_exe}"',
                f'"{version_txt}"',
                expected_sha256
            ]
        cmd_str = ' '.join(cmd)
        logger.info(f"[退出流程] 启动 updater 命令: {cmd_str}")
        subprocess.Popen(
            cmd_str,
            shell=True,
            cwd=str(APP_ROOT),
            close_fds=True,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS
        )
    except Exception as e:
        logger.error(f"[退出流程] 启动 updater 失败: {e}", exc_info=True)
        messagebox.showerror("更新错误", f"无法启动更新器：{e}")
    os._exit(0)

def check_and_update_on_startup() -> bool:
    try:
        current_ver = get_current_version()
        remote_ver, download_url, changelog, exe_sha256, remote_arch = get_remote_version()
        if not all([remote_ver, download_url, exe_sha256, remote_arch]):
            logger.warning("远程版本信息不完整，跳过更新")
            return True
        if parse_version(remote_ver) > parse_version(current_ver):
            should_update = show_update_dialog(current_ver, remote_ver, changelog)
            if should_update:
                perform_update_flow(current_ver, remote_ver, download_url, changelog, exe_sha256, remote_arch)
                return False
        return True
    except Exception as e:
        logger.error(f"启动更新检查异常: {e}", exc_info=True)
        return True

def run_main_gui():
    clean_old_mei_tempdirs()
    app = ctk.CTk()
    app.title("电商数据采集工具")
    app.geometry("600x400")
    app.resizable(False, False)
    ctk.CTkLabel(app, text="电商数据采集工具测试远程更新", font=("Microsoft YaHei", 20, "bold")).pack(pady=25)
    ctk.CTkLabel(app, text=f"当前版本：{get_current_version()}").pack()
    status_label = ctk.CTkLabel(app, text="状态：就绪", text_color="green")
    status_label.pack(pady=15)
    def start_crawling():
        status_label.configure(text="状态：爬取中...", text_color="blue")
        status_label.after(2000, lambda: status_label.configure(text="状态：完成", text_color="green"))
    ctk.CTkButton(app, text="开始采集", command=start_crawling, width=150, height=40).pack(pady=20)
    def on_close():
        clean_old_mei_tempdirs()
        app.destroy()
    app.protocol("WM_DELETE_WINDOW", on_close)
    app.mainloop()

if __name__ == '__main__':
    try:
        ctypes.windll.shcore.SetProcessDpiAwareness(1)
    except:
        pass
    ctk.set_appearance_mode("system")
    ctk.set_default_color_theme("blue")
    logger.info("========== 程序启动 ==========")
    if check_and_update_on_startup():
        run_main_gui()
    else:
        logger.info("更新流程已接管，主程序无需启动")
    logger.info("========== 程序结束 ==========")